package com.fliad.ai.core;

import org.noear.solon.ai.chat.ChatRole;
import org.noear.solon.ai.chat.ChatSessionDefault;
import org.noear.solon.ai.chat.message.ChatMessage;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.stream.Collectors;

public class ChatSessionImpl extends ChatSessionDefault {

    private static ChatMessage SYSTEM_MESSAGE = ChatMessage.ofSystem("You are a helpful assistant.");

    public ChatSessionImpl(String sessionId) {
        super(sessionId);
    }

    @Override
    public void addMessage(Collection<? extends ChatMessage> messages) {
        super.addMessage(messages.stream().filter(message -> ChatRole.SYSTEM != message.getRole()).collect(Collectors.toList()));
        messages.stream().filter(message -> ChatRole.SYSTEM == message.getRole()).findFirst()
                .ifPresent(message -> SYSTEM_MESSAGE = message);
    }

    @Override
    public List<ChatMessage> getMessages() {
        List<ChatMessage> userMessages = super.getMessages();
        List<ChatMessage> messages = null;
        if (userMessages.size() > 10) {
            messages = userMessages.subList(userMessages.size() - 10, userMessages.size());
        } else {
            messages = new ArrayList<>(userMessages);
        }
        messages.add(SYSTEM_MESSAGE);
        return messages;
    }

}
